import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import make_blobs

class KMeans:
    def __init__(self, k=3, max_iters=100, random_state=42):
        self.k = k
        self.max_iters = max_iters
        self.random_state = random_state
        self.centroids = None
        self.labels = None
        
    def initialize_centroids(self, X):
        """随机初始化质心"""
        np.random.seed(self.random_state)
        random_idx = np.random.permutation(X.shape[0])
        centroids = X[random_idx[:self.k]]
        return centroids
    
    def compute_distance(self, X, centroids):
        """计算每个点到质心的距离"""
        distances = np.zeros((X.shape[0], self.k))
        for i, centroid in enumerate(centroids):
            distances[:, i] = np.linalg.norm(X - centroid, axis=1)
        return distances
    
    def find_closest_centroid(self, distances):
        """找到最近的质心"""
        return np.argmin(distances, axis=1)
    
    def compute_centroids(self, X, labels):
        """重新计算质心"""
        centroids = np.zeros((self.k, X.shape[1]))
        for i in range(self.k):
            centroids[i] = np.mean(X[labels == i], axis=0)
        return centroids
    
    def fit(self, X):
        """训练K-means模型"""
        # 初始化质心
        self.centroids = self.initialize_centroids(X)
        
        for _ in range(self.max_iters):
            # 计算距离
            distances = self.compute_distance(X, self.centroids)
            
            # 分配标签
            self.labels = self.find_closest_centroid(distances)
            
            # 重新计算质心
            new_centroids = self.compute_centroids(X, self.labels)
            
            # 检查收敛
            if np.allclose(self.centroids, new_centroids):
                break
                
            self.centroids = new_centroids
            
        return self
    
    def predict(self, X):
        """预测新数据的类别"""
        distances = self.compute_distance(X, self.centroids)
        return self.find_closest_centroid(distances)

# 测试K-means算法
def test_kmeans():
    # 生成测试数据
    X, y_true = make_blobs(n_samples=300, centers=3, 
                          cluster_std=0.60, random_state=0)
    
    # 应用K-means
    kmeans = KMeans(k=3)
    kmeans.fit(X)
    y_pred = kmeans.labels
    
    # 可视化结果
    plt.figure(figsize=(12, 5))
    
    plt.subplot(1, 2, 1)
    plt.scatter(X[:, 0], X[:, 1], c=y_true, cmap='viridis')
    plt.title('真实标签')
    plt.colorbar()
    
    plt.subplot(1, 2, 2)
    plt.scatter(X[:, 0], X[:, 1], c=y_pred, cmap='viridis')
    plt.scatter(kmeans.centroids[:, 0], kmeans.centroids[:, 1], 
                marker='x', s=200, linewidths=3, color='red')
    plt.title('K-means聚类结果')
    plt.colorbar()
    
    plt.tight_layout()
    plt.show()
    
    return kmeans

# 运行K-means测试
kmeans_model = test_kmeans()